import json
import base64
import os
import sys
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.fernet import Fernet
from cryptography.hazmat.backends import default_backend


def generate_key(password: str, salt: bytes) -> bytes:
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100_000,
        backend=default_backend()
    )
    return base64.urlsafe_b64encode(kdf.derive(password.encode()))


def encrypt_json(json_path: str, password: str, output_path: str = None):
    with open(json_path, 'rb') as f:
        data = f.read()

    salt = os.urandom(16)
    key = generate_key(password, salt)
    fernet = Fernet(key)
    encrypted = fernet.encrypt(data)

    final_data = salt + encrypted  # 把 salt 附在开头
    if not output_path:
        output_path = json_path + ".enc"

    with open(output_path, 'wb') as f:
        f.write(final_data)


def decrypt_json(enc_path: str, password: str):
    with open(enc_path, 'rb') as f:
        content = f.read()

    salt = content[:16]
    encrypted = content[16:]

    key = generate_key(password, salt)
    fernet = Fernet(key)

    try:
        decrypted = fernet.decrypt(encrypted)
    except Exception as e:
        print("❌ 解密失败，密码可能错误或文件被篡改。")
        return

    try:
        return json.loads(decrypted.decode("utf-8")).get("usernames", [])
    except:
        return []


# 示例使用方法
if __name__ == "__main__":

    password = "<PASSWORD>"
    encrypt_json("../white_list.json", password)
    # decrypt_json(args.input, args.password)
